#!/usr/bin/python -i

from ceorModule import *

######################################
################ MAIN ################
## Is a 2-D box with lenght xL yL
if __name__ == "__main__":
  # And God says...
  N  = 2
  KT_eq = 0.1
  eps = 1.0
  sigma = 1.0
  delta_r_ini = 2.5*sigma
  dt = 0.001
  gas = Gas2D(10, 5, N, KT_eq, delta_r_ini)
  prefix="twoPart"
  ext = ".png"

  gas.particle[0].x_new = 1.0
  gas.particle[0].y_new = 2.5
  
  gas.particle[1].x_new = 3.5
  gas.particle[1].y_new = 2.5  
  
  gas.particle[0].vx_new = 0.0
  gas.particle[0].vy_new = 0.0
  
  gas.particle[1].vx_new = 0.0
  gas.particle[1].vy_new = 0.0
  
  
  #--- Plot initial positions
  gas.arrayPos()
  gas.arrayVel()
  print gas.X
  print gas.Y
  
  plt.figure(1)
  plt.grid()
  plt.xlabel('$ x $')
  plt.ylabel('$ y $')
  plt.scatter(gas.X, gas.Y, marker='s')
  plt.xlim((0,gas.xlen+1))
  plt.ylim((0,gas.ylen+1))
  # to be continued...  
  #outputfilename=prefix+"Fig01-XY_ini_end"+ext
  #plt.savefig(outputfilename)
  
  # 1. Calculamos F(t)
  gas.calculateForces(eps, sigma)
  gas.calculateEkin()
  print '-'*20
  print 'E_tot =', gas.Ekin+gas.Epot , 'Ekin_ini = ', gas.Ekin, 'Epot_ini = ', gas.Epot

  # ------ Begin SIMULATION ------
  total_time = 100000 #4900
  ttime_print = 1
  ttime = range(total_time)
  timeDim = total_time #/ttime_print + 1
  xt  = np.zeros([gas.N, timeDim])
  yt  = np.zeros([gas.N, timeDim])
  vxt = np.zeros([gas.N, timeDim])
  vyt = np.zeros([gas.N, timeDim])
  fxt = np.zeros([gas.N, timeDim])
  fyt = np.zeros([gas.N, timeDim])
  Ekint = np.zeros(timeDim)
  Epott = np.zeros(timeDim)
  Etott = np.zeros(timeDim)
  #######str(n).zfill(6)

  for i in ttime:
    gas.save_old() # fx_old=fx_new
    gas.stepPosition(dt) # input old variables = > output x_new
    gas.inBox()
    gas.calculateForces(eps, sigma)
    gas.stepVelocity(dt)
    # Save every each ttime_print
    #if i % ttime_print == 0 :
    i_print = i #/ttime_print
    gas.calculateEkin()
    Ekint[i_print] = gas.Ekin
    Epott[i_print] = gas.Epot
    #save particle trajectories
    for j in range(gas.N):
      xt[j][i_print]  = gas.particle[j].x_new
      yt[j][i_print]  = gas.particle[j].y_new
      vxt[j][i_print] = gas.particle[j].vx_new
      vyt[j][i_print] = gas.particle[j].vy_new
      fxt[j][i_print] = gas.particle[j].fx_new
      fyt[j][i_print] = gas.particle[j].fy_new
  # ------ End SIMULATION ------
  ttime = dt*np.array(ttime)
  Etott = Ekint+Epott
  gas.calculateEkin()
  print '-'*20
  print 'E_tot =', gas.Ekin+gas.Epot , 'Ekin_end = ', gas.Ekin, 'Epot_end = ', gas.Epot
  print '-- deltas --'*3
  print eps, dt, total_time, Etott.mean(), Etott.max() - Etott.min(),Etott.max(), Etott.min()
  
  # Save in file
  outdatafile=prefix+"Energy.dat"
  np.savetxt(outdatafile,zip(ttime,Epott, Ekint, Etott))
  
  gas.arrayPos()
  gas.arrayVel()
  plt.rc('text', usetex=True)
  plt.scatter(gas.X, gas.Y, marker='+')
  plt.xlim((0,gas.xlen+1))
  plt.ylim((0,gas.ylen+1))
  outputfilename=prefix+"Fig01-XY_ini_end"+ext
  plt.savefig(outputfilename)
  plt.clf()

  # TODO:
  plt.grid()
  plt.rc('font', size=25)
  plt.rc('text', usetex=True)
  plt.xlabel('$ t $')
  plt.ylabel('$ x $')
  plt.plot(ttime,xt[0])
  plt.plot(ttime,xt[1])
  plt.legend(["$x_0$", "$x_1$"], prop={'size':20})
  plt.ylim((-1,gas.xlen+1))
  outputfilename=prefix+"Fig03-TX_end"+ext
  plt.savefig(outputfilename)
  plt.clf()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ t $')
  plt.ylabel('$ v_x $')
  plt.plot(ttime,vxt[0])
  plt.plot(ttime,vxt[1])
  plt.legend(["$v_{x_0}$", "$v_{x_1}$"], prop={'size':20})
  plt.ylim((-2,2))
  outputfilename=prefix+"Fig04-TVx_end"+ext
  plt.savefig(outputfilename)
  plt.clf()
  # TODO:
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x $')
  plt.ylabel('$ v_x $')
  plt.plot(xt[0],vxt[0])
  plt.plot(xt[1],vxt[1])
  plt.legend(["$p_0$", "$p_1$"], prop={'size':20})
  outputfilename=prefix+"Fig05-XVx_end"+ext
  plt.savefig(outputfilename)
  plt.clf()
  # TODO:
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x_1 - x_0 $')
  plt.ylabel('$ E_{pot} $')
  plt.plot(xt[1]-xt[0],Epott)
  outputfilename=prefix+"Fig06-XEpot_end"+ext
  plt.savefig(outputfilename)
  plt.clf()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x_1 - x_0 $')
  plt.ylabel('$ E_{kin} $')
  plt.plot(xt[1]-xt[0],Ekint)
  outputfilename=prefix+"Fig07-XEkin_end"+ext
  plt.savefig(outputfilename)
  plt.clf()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x_1 - x_0 $')
  plt.ylabel('$ E_{tot} $')
  plt.plot(xt[1]-xt[0],Etott,'o-', markevery=5000)
  outputfilename=prefix+"Fig08-XEtot_end"+ext
  plt.savefig(outputfilename)
  plt.clf()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ t $')
  plt.plot(ttime,Epott)
  plt.plot(ttime,Ekint,'o-', markevery=5000)
  plt.plot(ttime,Etott)
  plt.legend(["$E_{pot}$","$E_{kin}$","$E_{tot}$"], prop={'size':20})
  plt.ylim((-1.0,1.0))
  outputfilename=prefix+"Fig09-tEtot_end"+ext
  plt.savefig(outputfilename)
  plt.clf()
  
  
  xLJ = np.arange(0.9,2.6,0.01)
  yLJ = 4*eps*((sigma/xLJ)**12 - (sigma/xLJ)**6)
  
  
  
"""
  # ------ Begin SIMULATION ------
  total_time = 100000 #4900
  ttime_print = 1
  ttime = range(total_time)
  timeDim = total_time #/ttime_print + 1
  xt  = np.zeros([gas.N, timeDim])
  yt  = np.zeros([gas.N, timeDim])
  vxt = np.zeros([gas.N, timeDim])
  vyt = np.zeros([gas.N, timeDim])
  fxt = np.zeros([gas.N, timeDim])
  fyt = np.zeros([gas.N, timeDim])
  Ekint = np.zeros(timeDim)
  Epott = np.zeros(timeDim)
  Etott = np.zeros(timeDim)
  #######str(n).zfill(6)

  for i in ttime:
    gas.save_old() # fx_old=fx_new
    gas.stepPosition(dt) # input old variables = > output x_new
    gas.inBox()
    gas.calculateForces(eps, sigma)
    gas.stepVelocity(dt)
    # Save every each ttime_print
    #if i % ttime_print == 0 :
    i_print = i #/ttime_print
    gas.calculateEkin()
    Ekint[i_print] = gas.Ekin
    Epott[i_print] = gas.Epot
    #save particle trajectories
    for j in range(gas.N):
      xt[j][i_print]  = gas.particle[j].x_new
      yt[j][i_print]  = gas.particle[j].y_new
      vxt[j][i_print] = gas.particle[j].vx_new
      vyt[j][i_print] = gas.particle[j].vy_new
      fxt[j][i_print] = gas.particle[j].fx_new
      fyt[j][i_print] = gas.particle[j].fy_new
  # ------ End SIMULATION ------
  ttime = dt*np.array(ttime)
  Etott = Ekint+Epott
  gas.calculateEkin()
  print '-'*20
  print 'E_tot =', gas.Ekin+gas.Epot , 'Ekin_end = ', gas.Ekin, 'Epot_end = ', gas.Epot
  print '-- deltas --'*3
  print eps, dt, total_time, Etott.mean(), Etott.max() - Etott.min(),Etott.max(), Etott.min()
  
  #plt.rc('text', usetex=True)
  plt.scatter(gas.X, gas.Y, marker='+')
  plt.xlim((0,gas.xlen+1))
  plt.ylim((0,gas.ylen+1))
  outputfilename=prefix+"Fig01-XY_ini_end"+ext
  plt.savefig(outputfilename)
  plt.close()

  # TODO:
  plt.grid()
  plt.rc('font', size=25)
  plt.rc('text', usetex=True)
  plt.xlabel('$ t $')
  plt.ylabel('$ x $')
  plt.plot(ttime,xt[0])
  plt.ylim((-1,gas.xlen+1))
  outputfilename=prefix+"Fig03-XT_end"+ext
  plt.savefig(outputfilename)
  plt.close()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ t $')
  plt.ylabel('$ v_x $')
  plt.plot(ttime,vxt[0])
  plt.ylim((-2,2))
  outputfilename=prefix+"Fig04-VxT_end"+ext
  plt.savefig(outputfilename)
  plt.close()
  # TODO:
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x $')
  plt.ylabel('$ v_x $')
  plt.plot(xt[0],vxt[0])
  outputfilename=prefix+"Fig05-XVx_end"+ext
  plt.savefig(outputfilename)
  plt.close()
  # TODO:
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x $')
  plt.ylabel('$ E_{pot} $')
  plt.plot(xt[0],Epott,'+')
  outputfilename=prefix+"Fig06-XEpot_end"+ext
  plt.savefig(outputfilename)
  plt.close()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x $')
  plt.ylabel('$ E_{kin} $')
  plt.plot(xt[0],Ekint)
  outputfilename=prefix+"Fig07-XEkin_end"+ext
  plt.savefig(outputfilename)
  plt.close()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ x $')
  plt.ylabel('$ E_{tot} $')
  plt.plot(xt[0],Etott,'o')
  outputfilename=prefix+"Fig08-XEtot_end"+ext
  plt.savefig(outputfilename)
  plt.close()
  # TODO: 
  plt.grid()
  plt.rc('text', usetex=True)
  plt.xlabel('$ t $')
  plt.plot(ttime,Epott)
  plt.plot(ttime,Ekint,'o', markevery=5000)
  plt.plot(ttime,Etott)
  plt.legend(["$E_{pot}$","$E_{kin}$","$E_{tot}$"], prop={'size':20})
  plt.ylim((-0.5,1.0))
  outputfilename=prefix+"Fig09-tEtot_end"+ext
  plt.savefig(outputfilename)
  plt.close()
  
"""
